{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f5f3a62d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as t\n",
    "import timm\n",
    "import torchvision as tv\n",
    "import torch.nn as nn\n",
    "from torch.optim import AdamW\n",
    "import wandb\n",
    "from data import construct_dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from models.vit_moe import ViTMoE\n",
    "from models.vit_shareparam import ViTMoEShareParam\n",
    "from engine import train_one_epoch, evalulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b83e2bee",
   "metadata": {},
   "outputs": [],
   "source": [
    "class args:\n",
    "    # data\n",
    "    data_root = \"./datasets\"\n",
    "    dataset_name = \"CIFAR10\"\n",
    "    batch_size = 256\n",
    "    # model\n",
    "    moe_model = ViTMoE\n",
    "    model_kwargs = dict(embed_dim=384, depth=8, num_experts=8)\n",
    "    # loss func\n",
    "    loss_func = nn.CrossEntropyLoss\n",
    "    # optimizer\n",
    "    lr = 5e-4\n",
    "    weight_decay = 0.05\n",
    "    # other\n",
    "    epoch = 300\n",
    "    device = \"cuda:0\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28bb278f",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set, test_set = construct_dataset(args.dataset_name, args.data_root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2ae2e919",
   "metadata": {},
   "outputs": [],
   "source": [
    "#train_set.targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1edf85cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#train_set.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "840ef82c",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(train_set.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c005c67",
   "metadata": {},
   "outputs": [],
   "source": [
    "idxx = train_set.targets\n",
    "idxx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b92f0e1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "argid = np.argsort(idxx)\n",
    "argid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "233a257f",
   "metadata": {},
   "outputs": [],
   "source": [
    "set(idxx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee2f3ded",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, Subset, TensorDataset\n",
    "id_dataset = Subset(train_set, indices=argid)\n",
    "for x in DataLoader(id_dataset):\n",
    "    print(x)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
